#!/usr/bin/env python
import itk
import sys
import os
import pylab as pb
import xseg.lungSeg as lungSeg


def vas_seg(fname, low_th, high_th, size, outputfile):
#Get segment the lung vasculature and grab the largest num of components
	reader = itk.ImageFileReader.ISS3.New()
        reader.SetFileName(fname)
        rsf = itk.RescaleIntensityImageFilter.ISS3ISS3.New()
        rsf.SetInput(reader.GetOutput())
        rsf.SetOutputMaximum(1000)
        rsf.SetOutputMinimum(0)
	btf = itk.BinaryThresholdImageFilter.ISS3ISS3.New()
	btf.SetInput(rsf.GetOutput())
	btf.SetLowerThreshold(low_th)
	btf.SetUpperThreshold(high_th)
	ccf = itk.ConnectedComponentImageFilter.ISS3IUL3.New()
        #ccf.SetFullyConnected(1)
	ccf.SetInput(btf.GetOutput())
        #reader2 = itk.ImageFileReader.ISS3.New()
        #reader2.SetFileName(fmask)
        #ccf.SetMaskImage(reader2.GetOutput())
	rcf = itk.RelabelComponentImageFilter.IUL3IUS3.New()
	rcf.SetInput(ccf.GetOutput())
        rcf.SetMinimumObjectSize(size)
        rcf.Update()
        #start = 1
        #end = 10
	#tf = itk.ThresholdImageFilter.IUS3.New()
	#tf.SetInput(rcf.GetOutput())
	#tf.SetLower(start)
	#tf.SetUpper(end)
	#tf.Update()
        writer= itk.ImageFileWriter.IUS3.New()
        writer.SetFileName(outputfile)
        writer.SetInput(rcf.GetOutput())
        writer.Update()
        return rcf.GetSizeOfObjectsInPixels()
def mask(fname, fmask) :
       reader1 = itk.ImageFileReader.ISS3.New()
       reader1.SetFileName(fname)
       reader2 = itk.ImageFileReader.ISS3.New()
       reader2.SetFileName(fmask)  
       mask = itk.MaskImageFilter.ISS3ISS3ISS3.New()
       mask.SetInput1(reader1.GetOutput())
       mask.SetInput2(reader2.GetOutput())
       writer = itk.ImageFileWriter.ISS3.New()
       tmp = os.path.splitext(fname)
       foutput = tmp[0] + "_masked.mha"
       writer.SetInput(mask.GetOutput())
       writer.SetFileName(foutput)
       writer.Update() 
       return foutput
   
def main():
       if len(sys.argv) < 6 :
            print " Usage : seg fname, low_th, high_th, size , mask"
            sys.exit(1)
       
       low_th = int(sys.argv[2])
       high_th = int(sys.argv[3])
       min_pixels = int(sys.argv[4]) 
       tmp = os.path.splitext(sys.argv[1])
       fmask = lungSeg.lungMaskExtraction(sys.argv[5])
       outputfile = tmp[0] +"__" + str(low_th) + "__"+str(high_th) + tmp[1] 
       foutput = mask(sys.argv[1], fmask)
       cmd = "rm " + fmask
       os.system(cmd)
       size = vas_seg(foutput, low_th, high_th, min_pixels , outputfile)
       print size
       pb.plot(size)
       pb.show() 
if __name__ == '__main__':
    main()

